import matplotlib.pyplot as plt
import os
os.environ['OMP_NUM_THREADS'] = '2'
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs

plt.rcParams['font.sans-serif'] = ['KaiTi']
plt.rcParams['axes.unicode_minus'] = False

# 1.生成数据集
X, y = make_blobs(n_samples=300, centers=4, cluster_std=2.5, random_state=42)

# 2. 画出散点图
fig, ax = plt.subplots(2, 1, figsize=(10, 8))
ax[0].scatter(X[:, 0], X[:, 1], s=50, label='原始数据')
ax[0].set_title('原始数据集')
ax[0].legend()

# 3.定义 KMeans 模型并训练
kmeans = KMeans(n_clusters=4)
kmeans.fit(X)

# 4. 获取聚类结果
centers = kmeans.cluster_centers_

# 5. 预测所有样本点分类标签
y_pred = kmeans.predict(X)

# 6. 画出聚类结果
ax[1].scatter(X[:, 0], X[:, 1], c=y_pred, s=50, label='聚类结果')
ax[1].scatter(centers[:, 0], centers[:, 1], c='r', s=200, label='簇中心')
ax[1].set_title('聚类结果')
ax[1].legend()
plt.show()
